#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netinet/in.h>
#include <string.h>
#include <sys/stat.h>

#define BUFFER_SIZE 1024    // 数据缓冲区大小
#define FILE_CHUNK_SIZE 512 // 文件分片大小
#define MAX_FILENAME 256    // 最大文件名长度

#define ERR_LOG(errmsg) do { \
    perror(errmsg); \
    exit(1); \
} while(0)

// 定义消息类型（与客户端保持一致）
typedef enum {
    MSG_TEXT = 1,        // 普通文本消息
    MSG_FILE_SEND,       // 发送文件请求
    MSG_FILE_RECV,       // 接收文件请求
    MSG_FILE_DATA,       // 文件数据包
    MSG_FILE_END,        // 文件传输结束
    MSG_ACK,             // 确认消息
    MSG_ERROR            // 错误消息
} msg_type_t;

// 消息头结构体
typedef struct {
    msg_type_t type;           // 消息类型
    int seq_num;               // 序列号
    int data_len;              // 数据长度
    char filename[MAX_FILENAME]; // 文件名
} msg_header_t;

// 完整消息结构体
typedef struct {
    msg_header_t header;
    char data[FILE_CHUNK_SIZE];
} message_t;

// 发送文件给客户端
int send_file_to_client(int sockfd, struct sockaddr_in *clientaddr, socklen_t addrlen, const char *filename) {
    FILE *fp = fopen(filename, "rb");
    if (fp == NULL) {
        // 发送错误消息
        message_t error_msg;
        error_msg.header.type = MSG_ERROR;
        error_msg.header.seq_num = 0;
        error_msg.header.data_len = 0;
        strncpy(error_msg.header.filename, filename, MAX_FILENAME - 1);
        error_msg.header.filename[MAX_FILENAME - 1] = '\0';
        
        //参数：套接字、要发送的内存地址、发送的数据的长度、0、客户端地址、地址长度
        sendto(sockfd, &error_msg, sizeof(msg_header_t), 0,
               (struct sockaddr *)clientaddr, addrlen);
        
        printf("文件不存在或无法读取: %s\n", filename);
        return -1;
    }

    // 获取文件大小
    fseek(fp, 0, SEEK_END);
    long file_size = ftell(fp);
    fseek(fp, 0, SEEK_SET);
    
    printf("开始发送文件给客户端 [%s:%d]: %s (大小: %ld 字节)\n",
           inet_ntoa(clientaddr->sin_addr), ntohs(clientaddr->sin_port), 
           filename, file_size);

    message_t msg;
    int seq_num = 0;
    int bytes_read;
    
    // 发送文件数据
    while ((bytes_read = fread(msg.data, 1, FILE_CHUNK_SIZE, fp)) > 0) {
        // 设置消息头
        msg.header.type = MSG_FILE_DATA;
        msg.header.seq_num = seq_num++;
        msg.header.data_len = bytes_read;
        strncpy(msg.header.filename, filename, MAX_FILENAME - 1);
        msg.header.filename[MAX_FILENAME - 1] = '\0';

        // 发送数据包
        if (sendto(sockfd, &msg, sizeof(msg_header_t) + bytes_read, 0,
                  (struct sockaddr *)clientaddr, addrlen) == -1) {
            perror("发送文件数据失败");
            fclose(fp);
            return -1;
        }

        // 等待ACK确认
        message_t ack_msg;
        if (recvfrom(sockfd, &ack_msg, sizeof(ack_msg), 0,
                    (struct sockaddr *)clientaddr, &addrlen) < 0) {
            perror("接收ACK失败");
            fclose(fp);
            return -1;
        }

        if (ack_msg.header.type != MSG_ACK || ack_msg.header.seq_num != seq_num - 1) {
            printf("收到错误的ACK\n");
            fclose(fp);
            return -1;
        }

        printf("已发送第 %d 个数据包给客户端 (%d 字节)\n", seq_num - 1, bytes_read);
    }

    // 发送文件结束标志
    msg.header.type = MSG_FILE_END;
    msg.header.seq_num = seq_num;
    msg.header.data_len = 0;
    strncpy(msg.header.filename, filename, MAX_FILENAME - 1);
    
    if (sendto(sockfd, &msg, sizeof(msg_header_t), 0,
              (struct sockaddr *)clientaddr, addrlen) == -1) {
        perror("发送文件结束标志失败");
        fclose(fp);
        return -1;
    }

    // 等待最后的ACK
    message_t final_ack;
    if (recvfrom(sockfd, &final_ack, sizeof(final_ack), 0,
                (struct sockaddr *)clientaddr, &addrlen) < 0) {
        perror("接收最终ACK失败");
    }

    fclose(fp);
    printf("文件发送完成给客户端: %s\n", filename);
    return 0;
}

// 从客户端接收文件
int receive_file_from_client(int sockfd, struct sockaddr_in *clientaddr, socklen_t addrlen, const char *filename) {
    // 创建服务器端文件名（添加前缀避免覆盖）
    char server_filename[MAX_FILENAME + 20];
    snprintf(server_filename, sizeof(server_filename), "server_%s", filename);
    
    FILE *fp = fopen(server_filename, "wb");
    if (fp == NULL) {
        printf("无法创建文件: %s\n", server_filename);
        return -1;
    }

    printf("开始从客户端 [%s:%d] 接收文件: %s -> %s\n",
           inet_ntoa(clientaddr->sin_addr), ntohs(clientaddr->sin_port),
           filename, server_filename);

    message_t msg, ack_msg;
    int expected_seq = 0;
    int total_bytes = 0;

    while (1) {
        // 接收数据包
        ssize_t recv_len = recvfrom(sockfd, &msg, sizeof(msg), 0,
                                   (struct sockaddr *)clientaddr, &addrlen);
        if (recv_len < 0) {
            perror("接收文件数据失败");
            fclose(fp);
            unlink(server_filename); // 删除不完整的文件
            return -1;
        }

        // 发送ACK
        ack_msg.header.type = MSG_ACK;
        ack_msg.header.seq_num = msg.header.seq_num;
        ack_msg.header.data_len = 0;
        
        if (sendto(sockfd, &ack_msg, sizeof(msg_header_t), 0,
                  (struct sockaddr *)clientaddr, addrlen) == -1) {
            perror("发送ACK失败");
        }

        // 处理不同类型的消息
        if (msg.header.type == MSG_FILE_DATA) {
            // 检查序列号
            if (msg.header.seq_num != expected_seq) {
                printf("收到乱序的数据包，期望: %d，实际: %d\n",
                       expected_seq, msg.header.seq_num);
                continue;
            }

            // 写入文件数据
            if (fwrite(msg.data, 1, msg.header.data_len, fp) != msg.header.data_len) {
                printf("写入文件失败\n");
                fclose(fp);
                unlink(server_filename);
                return -1;
            }

            total_bytes += msg.header.data_len;
            expected_seq++;
            printf("接收到第 %d 个数据包 (%d 字节)，总计: %d 字节\n",
                   msg.header.seq_num, msg.header.data_len, total_bytes);

        } else if (msg.header.type == MSG_FILE_END) {
            printf("文件接收完成: %s (总大小: %d 字节)\n", server_filename, total_bytes);
            break;
        }
    }

    fclose(fp);
    return 0;
}

int main(int argc, char const *argv[]) {
    if (argc < 3) {
        fprintf(stderr, "用法: %s <服务器IP> <端口>\n", argv[0]);
        exit(1);
    }

    // 第一步：创建套接字
    int sockfd;
    if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) == -1) {
        ERR_LOG("创建套接字失败");
    }

    // 设置套接字选项，允许地址重用
    int on = 1;
    if (setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0) {
        ERR_LOG("设置套接字选项失败");
    }

    // 第二步：绑定套接字到服务器地址
    struct sockaddr_in serveraddr;
    socklen_t addrlen = sizeof(serveraddr);

    serveraddr.sin_family = AF_INET;
    serveraddr.sin_addr.s_addr = inet_addr(argv[1]);
    serveraddr.sin_port = htons(atoi(argv[2]));

    if (bind(sockfd, (struct sockaddr *)&serveraddr, addrlen) == -1) {
        ERR_LOG("绑定失败");
    }

    printf("UDP文件传输服务器启动成功\n");
    printf("监听地址: %s:%s\n", argv[1], argv[2]);
    printf("当前工作目录: %s\n", getcwd(NULL, 0));
    printf("等待客户端连接...\n\n");

    // 主服务循环
    while (1) {
        message_t msg;
        struct sockaddr_in clientaddr;
        
        // 接收客户端消息
        ssize_t recv_len = recvfrom(sockfd, &msg, sizeof(msg), 0,
                                   (struct sockaddr *)&clientaddr, &addrlen);
        if (recv_len < 0) {
            perror("接收消息失败");
            continue;
        }

        printf("\n[客户端 %s:%d] ", 
               inet_ntoa(clientaddr.sin_addr), ntohs(clientaddr.sin_port));

        // 根据消息类型处理
        switch (msg.header.type) {
            case MSG_TEXT:
                // 处理普通文本消息
                msg.data[msg.header.data_len] = '\0'; // 确保字符串结尾
                printf("收到文本消息: %s\n", msg.data);
                
                // 回复消息（添加服务器标识）
                char reply[BUFFER_SIZE];
                snprintf(reply, sizeof(reply), "服务器收到: %s", msg.data);
                
                message_t reply_msg;
                reply_msg.header.type = MSG_TEXT;
                reply_msg.header.seq_num = 0;
                reply_msg.header.data_len = strlen(reply);
                strcpy(reply_msg.data, reply);
                
                if (sendto(sockfd, &reply_msg, sizeof(msg_header_t) + reply_msg.header.data_len, 0,
                          (struct sockaddr *)&clientaddr, addrlen) < 0) {
                    perror("发送回复失败");
                }
                break;

            case MSG_FILE_RECV:
                // 客户端请求下载文件
                printf("请求下载文件: %s\n", msg.header.filename);
                send_file_to_client(sockfd, &clientaddr, addrlen, msg.header.filename);
                break;

            case MSG_FILE_DATA:
                // 客户端发送文件数据，这是文件上传的第一个数据包
                printf("开始接收文件: %s\n", msg.header.filename);
                
                // 先处理第一个数据包
                char server_filename[MAX_FILENAME + 20];
                snprintf(server_filename, sizeof(server_filename), "server_%s", msg.header.filename);
                
                FILE *fp = fopen(server_filename, "wb");
                if (fp == NULL) {
                    printf("无法创建文件: %s\n", server_filename);
                    break;
                }

                // 写入第一个数据包
                if (fwrite(msg.data, 1, msg.header.data_len, fp) != msg.header.data_len) {
                    printf("写入文件失败\n");
                    fclose(fp);
                    unlink(server_filename);
                    break;
                }

                int total_bytes = msg.header.data_len;
                printf("接收到第 %d 个数据包 (%d 字节)，总计: %d 字节\n",
                       msg.header.seq_num, msg.header.data_len, total_bytes);

                // 发送第一个ACK
                message_t ack_msg;
                ack_msg.header.type = MSG_ACK;
                ack_msg.header.seq_num = msg.header.seq_num;
                ack_msg.header.data_len = 0;
                
                if (sendto(sockfd, &ack_msg, sizeof(msg_header_t), 0,
                          (struct sockaddr *)&clientaddr, addrlen) == -1) {
                    perror("发送ACK失败");
                }

                // 继续接收剩余数据包
                int expected_seq = msg.header.seq_num + 1;
                
                while (1) {
                    message_t next_msg;
                    ssize_t next_recv_len = recvfrom(sockfd, &next_msg, sizeof(next_msg), 0,
                                                    (struct sockaddr *)&clientaddr, &addrlen);
                    if (next_recv_len < 0) {
                        perror("接收文件数据失败");
                        break;
                    }

                    // 发送ACK
                    ack_msg.header.type = MSG_ACK;
                    ack_msg.header.seq_num = next_msg.header.seq_num;
                    ack_msg.header.data_len = 0;
                    
                    if (sendto(sockfd, &ack_msg, sizeof(msg_header_t), 0,
                              (struct sockaddr *)&clientaddr, addrlen) == -1) {
                        perror("发送ACK失败");
                    }

                    if (next_msg.header.type == MSG_FILE_DATA) {
                        // 检查序列号
                        if (next_msg.header.seq_num != expected_seq) {
                            printf("收到乱序的数据包，期望: %d，实际: %d\n",
                                   expected_seq, next_msg.header.seq_num);
                            continue;
                        }

                        // 写入文件数据
                        if (fwrite(next_msg.data, 1, next_msg.header.data_len, fp) != next_msg.header.data_len) {
                            printf("写入文件失败\n");
                            break;
                        }

                        total_bytes += next_msg.header.data_len;
                        expected_seq++;
                        printf("接收到第 %d 个数据包 (%d 字节)，总计: %d 字节\n",
                               next_msg.header.seq_num, next_msg.header.data_len, total_bytes);

                    } else if (next_msg.header.type == MSG_FILE_END) {
                        printf("文件接收完成: %s (总大小: %d 字节)\n", server_filename, total_bytes);
                        break;
                    }
                }

                fclose(fp);
                break;

            default:
                printf("收到未知类型的消息: %d\n", msg.header.type);
                break;
        }
    }

    close(sockfd);
    return 0;
}